#include "auto_navigation/planning/astar_planner.hpp"
#include <cmath>
#include <algorithm>

namespace auto_navigation {
namespace planning {

// 定义方向数组，8个方向
const int dx[8] = {-1, -1, 0, 1, 1, 1, 0, -1};
const int dy[8] = {0, 1, 1, 1, 0, -1, -1, -1};
const double cost[8] = {1.0, 1.414, 1.0, 1.414, 1.0, 1.414, 1.0, 1.414};

AStarPlanner::AStarPlanner() : PlannerInterface("astar"), map_received_(false), inflation_radius_(0.3) {
}

AStarPlanner::~AStarPlanner() {
}

bool AStarPlanner::initialize(ros::NodeHandle& nh) {
    // 从参数服务器获取参数
    nh.param<double>("astar/inflation_radius", inflation_radius_, 0.3);
    
    ROS_INFO("A* planner initialized with inflation radius: %.2f", inflation_radius_);
    return true;
}

bool AStarPlanner::setMap(const nav_msgs::OccupancyGrid& map) {
    map_ = map;
    map_received_ = true;
    
    // 获取地图属性
    width_ = map_.info.width;
    height_ = map_.info.height;
    resolution_ = map_.info.resolution;
    origin_x_ = map_.info.origin.position.x;
    origin_y_ = map_.info.origin.position.y;
    
    ROS_INFO("Map set for A* planner: %dx%d cells, resolution: %.2f", width_, height_, resolution_);
    return true;
}

bool AStarPlanner::plan(const geometry_msgs::PoseStamped& start,
                       const geometry_msgs::PoseStamped& goal,
                       nav_msgs::Path& path) {
    if (!map_received_) {
        ROS_ERROR("No map received yet");
        return false;
    }
    
    // 将起点和终点转换为网格坐标
    int start_x, start_y, goal_x, goal_y;
    if (!worldToGrid(start.pose.position.x, start.pose.position.y, start_x, start_y) ||
        !worldToGrid(goal.pose.position.x, goal.pose.position.y, goal_x, goal_y)) {
        ROS_ERROR("Start or goal position is outside the map");
        return false;
    }
    
    // 检查起点和终点是否有效
    if (!isValidCell(start_x, start_y) || !isValidCell(goal_x, goal_y)) {
        ROS_ERROR("Start or goal position is invalid");
        return false;
    }
    
    // 检查起点和终点是否被占用
    if (isOccupied(start_x, start_y)) {
        ROS_ERROR("Start position is occupied");
        return false;
    }
    
    if (isOccupied(goal_x, goal_y)) {
        ROS_ERROR("Goal position is occupied");
        return false;
    }
    
    // 初始化A*算法数据结构
    GridCell start_cell = {start_x, start_y};
    GridCell goal_cell = {goal_x, goal_y};
    
    // 优先队列，按f值排序
    auto cmp = [](const std::pair<double, GridCell>& a, const std::pair<double, GridCell>& b) {
        return a.first > b.first;
    };
    std::priority_queue<std::pair<double, GridCell>, 
                        std::vector<std::pair<double, GridCell>>, 
                        decltype(cmp)> open_set(cmp);
    
    // 记录已经访问过的节点
    std::unordered_map<GridCell, bool, GridCellHash> closed_set;
    
    // 记录路径
    std::unordered_map<GridCell, GridCell, GridCellHash> came_from;
    
    // 记录g值
    std::unordered_map<GridCell, double, GridCellHash> g_score;
    
    // 初始化起点
    g_score[start_cell] = 0;
    open_set.push(std::make_pair(heuristic(start_x, start_y, goal_x, goal_y), start_cell));
    
    // A*主循环
    while (!open_set.empty()) {
        // 获取f值最小的节点
        GridCell current = open_set.top().second;
        open_set.pop();
        
        // 如果到达终点，重建路径并返回
        if (current.x == goal_cell.x && current.y == goal_cell.y) {
            reconstructPath(came_from, current, path);
            return true;
        }
        
        // 如果已经访问过，跳过
        if (closed_set.find(current) != closed_set.end()) {
            continue;
        }
        
        // 标记为已访问
        closed_set[current] = true;
        
        // 遍历8个方向的邻居
        for (int i = 0; i < 8; ++i) {
            int nx = current.x + dx[i];
            int ny = current.y + dy[i];
            
            // 检查邻居是否有效
            if (!isValidCell(nx, ny) || isOccupied(nx, ny)) {
                continue;
            }
            
            GridCell neighbor = {nx, ny};
            
            // 如果已经访问过，跳过
            if (closed_set.find(neighbor) != closed_set.end()) {
                continue;
            }
            
            // 计算经过当前节点到达邻居的g值
            double tentative_g_score = g_score[current] + cost[i];
            
            // 如果邻居还没有g值，或者新的g值更小
            if (g_score.find(neighbor) == g_score.end() || tentative_g_score < g_score[neighbor]) {
                // 更新路径和g值
                came_from[neighbor] = current;
                g_score[neighbor] = tentative_g_score;
                
                // 计算f值并加入开放列表
                double f = tentative_g_score + heuristic(nx, ny, goal_x, goal_y);
                open_set.push(std::make_pair(f, neighbor));
            }
        }
    }
    
    // 如果无法到达终点
    ROS_ERROR("Could not find a path from start to goal");
    return false;
}

bool AStarPlanner::worldToGrid(double wx, double wy, int& gx, int& gy) const {
    gx = static_cast<int>((wx - origin_x_) / resolution_);
    gy = static_cast<int>((wy - origin_y_) / resolution_);
    
    return (gx >= 0 && gx < width_ && gy >= 0 && gy < height_);
}

void AStarPlanner::gridToWorld(int gx, int gy, double& wx, double& wy) const {
    wx = origin_x_ + (gx + 0.5) * resolution_;
    wy = origin_y_ + (gy + 0.5) * resolution_;
}

bool AStarPlanner::isValidCell(int x, int y) const {
    return (x >= 0 && x < width_ && y >= 0 && y < height_);
}

bool AStarPlanner::isOccupied(int x, int y) const {
    if (!isValidCell(x, y)) {
        return true;
    }
    
    int index = getIndex(x, y);
    return (map_.data[index] >= 50);  // 占用概率大于等于50%认为被占用
}

int AStarPlanner::getIndex(int x, int y) const {
    return y * width_ + x;
}

double AStarPlanner::heuristic(int x1, int y1, int x2, int y2) const {
    // 使用欧几里得距离作为启发函数
    return std::sqrt(std::pow(x2 - x1, 2) + std::pow(y2 - y1, 2));
}

void AStarPlanner::reconstructPath(const std::unordered_map<GridCell, GridCell, GridCellHash>& came_from,
                                  const GridCell& current,
                                  nav_msgs::Path& path) const {
    // 初始化路径
    path.header.frame_id = "map";
    path.header.stamp = ros::Time::now();
    path.poses.clear();
    
    // 从终点回溯到起点
    std::vector<GridCell> grid_path;
    GridCell cell = current;
    grid_path.push_back(cell);
    
    while (came_from.find(cell) != came_from.end()) {
        cell = came_from.at(cell);
        grid_path.push_back(cell);
    }
    
    // 反转路径，从起点到终点
    std::reverse(grid_path.begin(), grid_path.end());
    
    // 转换为世界坐标并添加到路径中
    for (const auto& grid_cell : grid_path) {
        geometry_msgs::PoseStamped pose;
        pose.header = path.header;
        
        double wx, wy;
        gridToWorld(grid_cell.x, grid_cell.y, wx, wy);
        
        pose.pose.position.x = wx;
        pose.pose.position.y = wy;
        pose.pose.position.z = 0.0;
        
        // 设置方向，简单地朝向下一个点
        pose.pose.orientation.w = 1.0;
        
        path.poses.push_back(pose);
    }
    
    ROS_INFO("Path found with %zu points", path.poses.size());
}

} // namespace planning
} // namespace auto_navigation 